# -*- coding: utf-8 -*-
import numpy as np
import random
import cv2
import matplotlib.pyplot as plt
import math
img_list=[]
label=[]
"""旋转裁切"""
degrees = 10
translate = .1
scale = .1
shear = 10
perspective=0.001
im = cv2.imread("./data/images/bus.jpg")
im = cv2.cvtColor(im,cv2.COLOR_BGR2RGB)
img_list.append(im)
label.append("原图")
height = im.shape[0]
width = im.shape[1]

"""向右下平移"""
# Translation
T = np.eye(3)  # 对角矩阵
# 图片的宽和高做处理
T[0, 2] = random.uniform(0.5 - translate, 0.5 + translate) * width  # x translation (pixels)
T[1, 2] = random.uniform(0.5 - translate, 0.5 + translate) * height  # y translation (pixels)
print("T",T)
im2 = cv2.warpPerspective(im, T, dsize=(width, height), borderValue=(114, 114, 114))
img_list.append(im2)
label.append("右下平移")

"""Shear   错切、非垂直投影"""
S = np.eye(3)
S[0, 1] = math.tan(random.uniform(-shear, shear) * math.pi / 180)  # x shear (deg)
S[1, 0] = math.tan(random.uniform(-shear, shear) * math.pi / 180)  # y shear (deg)
print("S",S)
im3 = cv2.warpPerspective(im, S, dsize=(width, height), borderValue=(114, 114, 114))
img_list.append(im3)
label.append("投影")

"""Rotation and Scale 旋转裁切"""
R = np.eye(3)  # 对角矩阵
a = random.uniform(-degrees, degrees)
# a += random.choice([-180, -90, 0, 90])  # add 90deg rotations to small rotations
s = random.uniform(1 - scale, 1 + scale)
# s = 2 ** random.uniform(-scale, scale)
R[:2] = cv2.getRotationMatrix2D(angle=a, center=(0, 0), scale=s)
print("R",R)
im1 = cv2.warpPerspective(im, R, dsize=(width, height), borderValue=(114, 114, 114))
img_list.append(im1)
label.append("旋转裁切")

"""Perspective 透视变换"""
P = np.eye(3)  # 构造对角矩阵
P[2, 0] = random.uniform(-perspective, perspective)  # x perspective (about y)
P[2, 1] = random.uniform(-perspective, perspective)  # y perspective (about x)
print("P",P)
im4 = cv2.warpPerspective(im, P, dsize=(width, height), borderValue=(114, 114, 114))
img_list.append(im4)
label.append("透视变换")

""" # Center  """
C = np.eye(3)
C[0, 2] = -im.shape[1] / 2  # x translation (pixels)
C[1, 2] = -im.shape[0] / 2  # y translation (pixels)
print("C",C)
im5 = cv2.warpPerspective(im, C, dsize=(width, height), borderValue=(114, 114, 114))
img_list.append(im5)
label.append("左上平移")

"""马赛克拼接"""
labels4, segments4 = [], []
s = img_size = 640
mosaic_border = [-img_size // 2, -img_size // 2]
yc, xc = (int(random.uniform(-x, 2 * s + x)) for x in mosaic_border)  # mosaic center x, y
# indices = [index] + random.choices(self.indices, k=3)  # 3 additional image indices
# random.shuffle(indices)
indices ={
    "./data/images/bus.jpg",
    "./data/images/zidane.jpg",
    "./data/images/zidane.jpg",
    "./data/images/bus.jpg"
}
img4 = None
for i, index in enumerate(indices):   # 枚举都是从字典中走
    # # Load image
    # img, _, (h, w) = load_image(index)
    img = cv2.imread(index)
    img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
    h,w,_ = np.shape(img)
    img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8)  # base image with 4 tiles
    # place img in img4
    if i == 0:  # top left
        x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc  # xmin, ymin, xmax, ymax (large image)
        x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h  # xmin, ymin, xmax, ymax (small image)
    elif i == 1:  # top right
        x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc
        x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h
    elif i == 2:  # bottom left
        x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)
        x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h)
    elif i == 3:  # bottom right
        x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)
        x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)

    img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b]  # img4[ymin:ymax, xmin:xmax]
img_list.append(img4[:,:,::-1])
label.append("马赛克拼接")

"""mixup图像融合"""
im2 = cv2.imread("./data/images/zidane.jpg")
im2 = cv2.cvtColor(im2,cv2.COLOR_BGR2RGB)
im = cv2.resize(im,(640,640))
im2 = cv2.resize(im2,(640,640))
# Applies MixUp augmentation https://arxiv.org/pdf/1710.09412.pdf
r = np.random.beta(32.0, 32.0)  # mixup ratio, alpha=beta=32.0
im3 = (im * r + im2 * (1 - r)).astype(np.uint8)
img_list.append(im3)
label.append("mixup")

# 可视化操作
plt.figure()
plt.rcParams['font.sans-serif']='SimHei'
for i in range(len(img_list)):
    plt.subplot(4,4,i+1)
    plt.title(label[i])
    plt.imshow(img_list[i])
    plt.tight_layout()
plt.show()


